import torch
import torch.nn as nn
from transformers import BertModel

class BertSpamClassifier(nn.Module):
    def __init__(self, pretrained_model_name='bert-base-uncased', num_labels=2):
        super().__init__()
        self.bert = BertModel.from_pretrained(pretrained_model_name)
        self.dropout = nn.Dropout(0.3)
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)
        torch.device("cuda")
        if (torch.cuda.is_available()):
            self.device = torch.device('cuda')
        #  增加对mac M系列处理器版本的支持
        elif (torch.mps.is_available()):
            self.device = torch.device('mps')
        else:
            self.device = torch.device('cpu')
        self.to(self.device)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        x = self.dropout(pooled_output)
        logits = self.classifier(x)
        return logits 